{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rdygM0QvGBVz"
   },
   "source": [
    "# 基于自注意力机制的翻译系统"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "K0REVxb6GBV0"
   },
   "source": [
    "## 1. 数据处理\n",
    "- 读取数据\n",
    "- 分别保存为inputs，outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CMbb7XRqGBV1"
   },
   "outputs": [],
   "source": [
    "with open('cmn.txt', 'r', encoding='utf8') as f:\n",
    "    data = f.readlines()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "AIZ7VTlIGBV6",
    "outputId": "66810816-dae8-49c4-d0f0-9393fca6ce1a"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [00:00<00:00, 473991.57it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "inputs = []\n",
    "outputs = []\n",
    "for line in tqdm(data[:10000]):\n",
    "    [en, ch] = line.strip('\\n').split('\\t')\n",
    "    inputs.append(en.replace(',',' ,')[:-1].lower())\n",
    "    outputs.append(ch[:-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "q8s8l1tDGBV_",
    "outputId": "8b31a7b6-7237-4a52-8522-76d26abcb2b0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['hi', 'hi', 'run', 'wait', 'hello', 'i try', 'i won', 'oh no', 'cheers', 'he ran']\n"
     ]
    }
   ],
   "source": [
    "print(inputs[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "kFGSTmYsGBWC",
    "outputId": "1bb178c2-0cef-48ba-c527-7d0c27288d64",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['嗨', '你好', '你用跑的', '等等', '你好', '让我来', '我赢了', '不会吧', '乾杯', '他跑了']\n"
     ]
    }
   ],
   "source": [
    "print(outputs[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vsyUGzxYGBWG"
   },
   "source": [
    "### 1.1 英文分词\n",
    "我们将英文用空格隔开即可，但是需要稍微修改一下，将大写字母全部用小写字母代替。在上文中使用`.lower`进行了替代。\n",
    "\n",
    "```py\n",
    "for line in tqdm(data):\n",
    "    [en, ch] = line.strip('\\n').split('\\t')\n",
    "    inputs.append(en[:-1].lower())\n",
    "    outputs.append(ch[:-1])\n",
    "\n",
    "```\n",
    "此处我们只需要将英文用空格分开即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Diqml7onGBWG"
   },
   "outputs": [],
   "source": [
    "inputs = [en.split(' ') for en in inputs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "r35hryRpGBWL",
    "outputId": "62fd64c7-36b7-4964-9e15-0fe68deb7e39"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['hi'], ['hi'], ['run'], ['wait'], ['hello'], ['i', 'try'], ['i', 'won'], ['oh', 'no'], ['cheers'], ['he', 'ran']]\n"
     ]
    }
   ],
   "source": [
    "print(inputs[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jenFGFO0GBWN"
   },
   "source": [
    "### 1.2 中文分词\n",
    "- 中文分词选择结巴分词工具。\n",
    "```py\n",
    "import jieba\n",
    "outputs = [[char for char in jieba.cut(line) if char != ' '] for line in outputs]\n",
    "```\n",
    "- 也可以用hanlp。\n",
    "```py\n",
    "from pyhanlp import *\n",
    "outputs = [[term.word for term in HanLP.segment(line) if term.word != ' '] for line in outputs]\n",
    "```\n",
    "- 或者按字分词？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 55
    },
    "colab_type": "code",
    "id": "a5r0uPPPGBWO",
    "outputId": "37a29150-ea61-4b94-8b40-29b9e3e2adbd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['你', '不應', '該', '去', '那裡', '的'], ['你', '以前', '吸煙', '，', '不是', '嗎'], ['你現', '在', '最好', '回家'], ['你', '今天', '最好', '不要', '出門'], ['你', '滑雪', '比', '我', '好'], ['你', '正在', '把', '我', '杯子', '里', '的', '东西', '喝掉'], ['你', '并', '不', '满意', '，', '对', '吧'], ['你', '病', '了', '，', '该', '休息', '了'], ['你', '很', '勇敢', '，', '不是', '嗎'], ['你', '的', '意志力', '很強']]\n"
     ]
    }
   ],
   "source": [
    "import jieba\n",
    "jieba_outputs = [[char for char in jieba.cut(line) if char != ' '] for line in outputs[-10:]]\n",
    "print(jieba_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "9f5gw9AMGBWZ",
    "outputId": "966a121b-3fae-40b3-b07d-d18a477b1159"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [00:00<00:00, 11981.68it/s]\n"
     ]
    }
   ],
   "source": [
    "outputs = [[char for char in jieba.cut(line) if char != ' '] for line in tqdm(outputs)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fp-1z59YGBWd"
   },
   "source": [
    "### 1.3 生成字典\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "XPDmTQSFGBWe",
    "outputId": "e1f97b9e-a753-4327-c3d2-0542dbcc5ebd"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [00:00<00:00, 20585.73it/s]\n",
      "100%|██████████| 10000/10000 [00:01<00:00, 7808.17it/s]\n"
     ]
    }
   ],
   "source": [
    "def get_vocab(data, init=['<PAD>']):\n",
    "    vocab = init\n",
    "    for line in tqdm(data):\n",
    "        for word in line:\n",
    "            if word not in vocab:\n",
    "                vocab.append(word)\n",
    "    return vocab\n",
    "\n",
    "SOURCE_CODES = ['<PAD>']\n",
    "TARGET_CODES = ['<PAD>', '<GO>', '<EOS>']\n",
    "encoder_vocab = get_vocab(inputs, init=SOURCE_CODES)\n",
    "decoder_vocab = get_vocab(outputs, init=TARGET_CODES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "IFhtgzuuGBWj",
    "outputId": "5d5c459b-8814-4cfe-aa8f-a065aad9c390"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<PAD>', 'hi', 'run', 'wait', 'hello', 'i', 'try', 'won', 'oh', 'no']\n",
      "['<PAD>', '<GO>', '<EOS>', '嗨', '你好', '你', '用', '跑', '的', '等等']\n"
     ]
    }
   ],
   "source": [
    "print(encoder_vocab[:10])\n",
    "print(decoder_vocab[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5UBCwOCKGBWl"
   },
   "source": [
    "### 1.4 数据生成器\n",
    "\n",
    "翻译系统解码器端需要输入和输出，因此解码器需要拼凑出输入序列和输出序列。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sRhUa227GBWm"
   },
   "outputs": [],
   "source": [
    "encoder_inputs = [[encoder_vocab.index(word) for word in line] for line in inputs]\n",
    "decoder_inputs = [[decoder_vocab.index('<GO>')] + [decoder_vocab.index(word) for word in line] for line in outputs]\n",
    "decoder_targets = [[decoder_vocab.index(word) for word in line] + [decoder_vocab.index('<EOS>')] for line in outputs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "dxzUF9R0GBWp",
    "outputId": "8d736a7f-de4f-4184-8693-d5343dcc1c5f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1, 3], [1, 4], [1, 5, 6, 7, 8], [1, 9]]\n",
      "[[3, 2], [4, 2], [5, 6, 7, 8, 2], [9, 2]]\n"
     ]
    }
   ],
   "source": [
    "print(decoder_inputs[:4])\n",
    "print(decoder_targets[:4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ITFlqyeIGBWu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def get_batch(encoder_inputs, decoder_inputs, decoder_targets, batch_size=4):\n",
    "    batch_num = len(encoder_inputs) // batch_size\n",
    "    for k in range(batch_num):\n",
    "        begin = k * batch_size\n",
    "        end = begin + batch_size\n",
    "        en_input_batch = encoder_inputs[begin:end]\n",
    "        de_input_batch = decoder_inputs[begin:end]\n",
    "        de_target_batch = decoder_targets[begin:end]\n",
    "        max_en_len = max([len(line) for line in en_input_batch])\n",
    "        max_de_len = max([len(line) for line in de_input_batch])\n",
    "        en_input_batch = np.array([line + [0] * (max_en_len-len(line)) for line in en_input_batch])\n",
    "        de_input_batch = np.array([line + [0] * (max_de_len-len(line)) for line in de_input_batch])\n",
    "        de_target_batch = np.array([line + [0] * (max_de_len-len(line)) for line in de_target_batch])\n",
    "        yield en_input_batch, de_input_batch, de_target_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 197
    },
    "colab_type": "code",
    "id": "ad92b5GsGBWz",
    "outputId": "48f623da-e04d-4f70-b386-f19ba5cf0a45"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[1],\n",
       "        [1],\n",
       "        [2],\n",
       "        [3]]), array([[1, 3, 0, 0, 0],\n",
       "        [1, 4, 0, 0, 0],\n",
       "        [1, 5, 6, 7, 8],\n",
       "        [1, 9, 0, 0, 0]]), array([[3, 2, 0, 0, 0],\n",
       "        [4, 2, 0, 0, 0],\n",
       "        [5, 6, 7, 8, 2],\n",
       "        [9, 2, 0, 0, 0]]))"
      ]
     },
     "execution_count": 74,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch = get_batch(encoder_inputs, decoder_inputs, decoder_targets, batch_size=4)\n",
    "next(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_FDZYnOuGBW3"
   },
   "source": [
    "## 2. 构建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zgxTCsGdGBW3"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IlpHtw43GBW7"
   },
   "source": [
    "### 2.1 构造建模组件\n",
    "下面代码实现了图片结构中的各个功能组件。\n",
    "![image.png](attachment:image.png)\n",
    "\n",
    "#### layer norm层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CUhfqeTXGBW8"
   },
   "outputs": [],
   "source": [
    "def normalize(inputs, \n",
    "              epsilon = 1e-8,\n",
    "              scope=\"ln\",\n",
    "              reuse=None):\n",
    "    '''Applies layer normalization.\n",
    "\n",
    "    Args:\n",
    "      inputs: A tensor with 2 or more dimensions, where the first dimension has\n",
    "        `batch_size`.\n",
    "      epsilon: A floating number. A very small number for preventing ZeroDivision Error.\n",
    "      scope: Optional scope for `variable_scope`.\n",
    "      reuse: Boolean, whether to reuse the weights of a previous layer\n",
    "        by the same name.\n",
    "\n",
    "    Returns:\n",
    "      A tensor with the same shape and data dtype as `inputs`.\n",
    "    '''\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "        inputs_shape = inputs.get_shape()\n",
    "        params_shape = inputs_shape[-1:]\n",
    "\n",
    "        mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)\n",
    "        beta= tf.Variable(tf.zeros(params_shape))\n",
    "        gamma = tf.Variable(tf.ones(params_shape))\n",
    "        normalized = (inputs - mean) / ( (variance + epsilon) ** (.5) )\n",
    "        outputs = gamma * normalized + beta\n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3ptiChtEGBW-"
   },
   "source": [
    "#### embedding层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kATHZ4V0GBW-"
   },
   "outputs": [],
   "source": [
    "def embedding(inputs, \n",
    "              vocab_size, \n",
    "              num_units, \n",
    "              zero_pad=True, \n",
    "              scale=True,\n",
    "              scope=\"embedding\", \n",
    "              reuse=None):\n",
    "    '''Embeds a given tensor.\n",
    "    Args:\n",
    "      inputs: A `Tensor` with type `int32` or `int64` containing the ids\n",
    "         to be looked up in `lookup table`.\n",
    "      vocab_size: An int. Vocabulary size.\n",
    "      num_units: An int. Number of embedding hidden units.\n",
    "      zero_pad: A boolean. If True, all the values of the fist row (id 0)\n",
    "        should be constant zeros.\n",
    "      scale: A boolean. If True. the outputs is multiplied by sqrt num_units.\n",
    "      scope: Optional scope for `variable_scope`.\n",
    "      reuse: Boolean, whether to reuse the weights of a previous layer\n",
    "        by the same name.\n",
    "    Returns:\n",
    "      A `Tensor` with one more rank than inputs's. The last dimensionality\n",
    "        should be `num_units`.\n",
    "\n",
    "    For example,\n",
    "\n",
    "    ```\n",
    "    import tensorflow as tf\n",
    "\n",
    "    inputs = tf.to_int32(tf.reshape(tf.range(2*3), (2, 3)))\n",
    "    outputs = embedding(inputs, 6, 2, zero_pad=True)\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        print sess.run(outputs)\n",
    "    >>\n",
    "    [[[ 0.          0.        ]\n",
    "      [ 0.09754146  0.67385566]\n",
    "      [ 0.37864095 -0.35689294]]\n",
    "     [[-1.01329422 -1.09939694]\n",
    "      [ 0.7521342   0.38203377]\n",
    "      [-0.04973143 -0.06210355]]]\n",
    "    ```\n",
    "\n",
    "    ```\n",
    "    import tensorflow as tf\n",
    "\n",
    "    inputs = tf.to_int32(tf.reshape(tf.range(2*3), (2, 3)))\n",
    "    outputs = embedding(inputs, 6, 2, zero_pad=False)\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        print sess.run(outputs)\n",
    "    >>\n",
    "    [[[-0.19172323 -0.39159766]\n",
    "      [-0.43212751 -0.66207761]\n",
    "      [ 1.03452027 -0.26704335]]\n",
    "     [[-0.11634696 -0.35983452]\n",
    "      [ 0.50208133  0.53509563]\n",
    "      [ 1.22204471 -0.96587461]]]\n",
    "    ```    \n",
    "    '''\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "        lookup_table = tf.get_variable('lookup_table',\n",
    "                                       dtype=tf.float32,\n",
    "                                       shape=[vocab_size, num_units],\n",
    "                                       initializer=tf.contrib.layers.xavier_initializer())\n",
    "        if zero_pad:\n",
    "            lookup_table = tf.concat((tf.zeros(shape=[1, num_units]),\n",
    "                                      lookup_table[1:, :]), 0)\n",
    "        outputs = tf.nn.embedding_lookup(lookup_table, inputs)\n",
    "\n",
    "        if scale:\n",
    "            outputs = outputs * (num_units ** 0.5) \n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IffFl8nhGBXB"
   },
   "source": [
    "#### multihead层\n",
    "该层实现了下面功能：\n",
    "![image.png](attachment:image.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "P4dVC9OaGBXC"
   },
   "outputs": [],
   "source": [
    "def multihead_attention(key_emb,\n",
    "                        que_emb,\n",
    "                        queries, \n",
    "                        keys, \n",
    "                        num_units=None, \n",
    "                        num_heads=8, \n",
    "                        dropout_rate=0,\n",
    "                        is_training=True,\n",
    "                        causality=False,\n",
    "                        scope=\"multihead_attention\", \n",
    "                        reuse=None):\n",
    "    '''Applies multihead attention.\n",
    "    \n",
    "    Args:\n",
    "      queries: A 3d tensor with shape of [N, T_q, C_q].\n",
    "      keys: A 3d tensor with shape of [N, T_k, C_k].\n",
    "      num_units: A scalar. Attention size.\n",
    "      dropout_rate: A floating point number.\n",
    "      is_training: Boolean. Controller of mechanism for dropout.\n",
    "      causality: Boolean. If true, units that reference the future are masked. \n",
    "      num_heads: An int. Number of heads.\n",
    "      scope: Optional scope for `variable_scope`.\n",
    "      reuse: Boolean, whether to reuse the weights of a previous layer\n",
    "        by the same name.\n",
    "        \n",
    "    Returns\n",
    "      A 3d tensor with shape of (N, T_q, C)  \n",
    "    '''\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "        # Set the fall back option for num_units\n",
    "        if num_units is None:\n",
    "            num_units = queries.get_shape().as_list[-1]\n",
    "        \n",
    "        # Linear projections\n",
    "        Q = tf.layers.dense(queries, num_units, activation=tf.nn.relu) # (N, T_q, C)\n",
    "        K = tf.layers.dense(keys, num_units, activation=tf.nn.relu) # (N, T_k, C)\n",
    "        V = tf.layers.dense(keys, num_units, activation=tf.nn.relu) # (N, T_k, C)\n",
    "        \n",
    "        # Split and concat\n",
    "        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0) # (h*N, T_q, C/h) \n",
    "        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0) # (h*N, T_k, C/h) \n",
    "        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0) # (h*N, T_k, C/h) \n",
    "\n",
    "        # Multiplication\n",
    "        outputs = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1])) # (h*N, T_q, T_k)\n",
    "        \n",
    "        # Scale\n",
    "        outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)\n",
    "        \n",
    "        # Key Masking\n",
    "        key_masks = tf.sign(tf.abs(tf.reduce_sum(key_emb, axis=-1))) # (N, T_k)\n",
    "        key_masks = tf.tile(key_masks, [num_heads, 1]) # (h*N, T_k)\n",
    "        key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1]) # (h*N, T_q, T_k)\n",
    "        \n",
    "        paddings = tf.ones_like(outputs)*(-2**32+1)\n",
    "        outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs) # (h*N, T_q, T_k)\n",
    "  \n",
    "        # Causality = Future blinding\n",
    "        if causality:\n",
    "            diag_vals = tf.ones_like(outputs[0, :, :]) # (T_q, T_k)\n",
    "            tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense() # (T_q, T_k)\n",
    "            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1]) # (h*N, T_q, T_k)\n",
    "   \n",
    "            paddings = tf.ones_like(masks)*(-2**32+1)\n",
    "            outputs = tf.where(tf.equal(masks, 0), paddings, outputs) # (h*N, T_q, T_k)\n",
    "  \n",
    "        # Activation\n",
    "        outputs = tf.nn.softmax(outputs) # (h*N, T_q, T_k)\n",
    "         \n",
    "        # Query Masking\n",
    "        query_masks = tf.sign(tf.abs(tf.reduce_sum(que_emb, axis=-1))) # (N, T_q)\n",
    "        query_masks = tf.tile(query_masks, [num_heads, 1]) # (h*N, T_q)\n",
    "        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]]) # (h*N, T_q, T_k)\n",
    "        outputs *= query_masks # broadcasting. (N, T_q, C)\n",
    "          \n",
    "        # Dropouts\n",
    "        outputs = tf.layers.dropout(outputs, rate=dropout_rate, training=tf.convert_to_tensor(is_training))\n",
    "               \n",
    "        # Weighted sum\n",
    "        outputs = tf.matmul(outputs, V_) # ( h*N, T_q, C/h)\n",
    "        \n",
    "        # Restore shape\n",
    "        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2 ) # (N, T_q, C)\n",
    "              \n",
    "        # Residual connection\n",
    "        outputs += queries\n",
    "              \n",
    "        # Normalize\n",
    "        outputs = normalize(outputs) # (N, T_q, C)\n",
    " \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "US_RmTCMGBXE"
   },
   "source": [
    "#### feedforward\n",
    "\n",
    "两层全连接，用卷积模拟加速运算，也可以使用dense层。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JFLKUuUmGBXF"
   },
   "outputs": [],
   "source": [
    "def feedforward(inputs, \n",
    "                num_units=[2048, 512],\n",
    "                scope=\"multihead_attention\", \n",
    "                reuse=None):\n",
    "    '''Point-wise feed forward net.\n",
    "    \n",
    "    Args:\n",
    "      inputs: A 3d tensor with shape of [N, T, C].\n",
    "      num_units: A list of two integers.\n",
    "      scope: Optional scope for `variable_scope`.\n",
    "      reuse: Boolean, whether to reuse the weights of a previous layer\n",
    "        by the same name.\n",
    "        \n",
    "    Returns:\n",
    "      A 3d tensor with the same shape and dtype as inputs\n",
    "    '''\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "        # Inner layer\n",
    "        params = {\"inputs\": inputs, \"filters\": num_units[0], \"kernel_size\": 1,\n",
    "                  \"activation\": tf.nn.relu, \"use_bias\": True}\n",
    "        outputs = tf.layers.conv1d(**params)\n",
    "        \n",
    "        # Readout layer\n",
    "        params = {\"inputs\": outputs, \"filters\": num_units[1], \"kernel_size\": 1,\n",
    "                  \"activation\": None, \"use_bias\": True}\n",
    "        outputs = tf.layers.conv1d(**params)\n",
    "        \n",
    "        # Residual connection\n",
    "        outputs += inputs\n",
    "        \n",
    "        # Normalize\n",
    "        outputs = normalize(outputs)\n",
    "    \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QIbOoVMFGBXH"
   },
   "source": [
    "#### label_smoothing.\n",
    "对于训练有好处，将0变为接近零的小数，1变为接近1的数，原文：\n",
    "\n",
    "During training, we employed label smoothing of value \u000fls = 0.1 [36]. This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3Su_os4AGBXI"
   },
   "outputs": [],
   "source": [
    "def label_smoothing(inputs, epsilon=0.1):\n",
    "    '''Applies label smoothing. See https://arxiv.org/abs/1512.00567.\n",
    "    \n",
    "    Args:\n",
    "      inputs: A 3d tensor with shape of [N, T, V], where V is the number of vocabulary.\n",
    "      epsilon: Smoothing rate.\n",
    "    \n",
    "    For example,\n",
    "    \n",
    "    ```\n",
    "    import tensorflow as tf\n",
    "    inputs = tf.convert_to_tensor([[[0, 0, 1], \n",
    "       [0, 1, 0],\n",
    "       [1, 0, 0]],\n",
    "      [[1, 0, 0],\n",
    "       [1, 0, 0],\n",
    "       [0, 1, 0]]], tf.float32)\n",
    "       \n",
    "    outputs = label_smoothing(inputs)\n",
    "    \n",
    "    with tf.Session() as sess:\n",
    "        print(sess.run([outputs]))\n",
    "    \n",
    "    >>\n",
    "    [array([[[ 0.03333334,  0.03333334,  0.93333334],\n",
    "        [ 0.03333334,  0.93333334,  0.03333334],\n",
    "        [ 0.93333334,  0.03333334,  0.03333334]],\n",
    "       [[ 0.93333334,  0.03333334,  0.03333334],\n",
    "        [ 0.93333334,  0.03333334,  0.03333334],\n",
    "        [ 0.03333334,  0.93333334,  0.03333334]]], dtype=float32)]   \n",
    "    ```    \n",
    "    '''\n",
    "    K = inputs.get_shape().as_list()[-1] # number of channels\n",
    "    return ((1-epsilon) * inputs) + (epsilon / K)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iwSIPvYtGBXK"
   },
   "source": [
    "### 2.2 搭建模型\n",
    "模型实现下图结构：\n",
    "![image.png](attachment:image.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oAkQDzSGGBXK"
   },
   "outputs": [],
   "source": [
    "class Graph():\n",
    "    def __init__(self, is_training=True):\n",
    "        tf.reset_default_graph()\n",
    "        self.is_training = arg.is_training\n",
    "        self.hidden_units = arg.hidden_units\n",
    "        self.input_vocab_size = arg.input_vocab_size\n",
    "        self.label_vocab_size = arg.label_vocab_size\n",
    "        self.num_heads = arg.num_heads\n",
    "        self.num_blocks = arg.num_blocks\n",
    "        self.max_length = arg.max_length\n",
    "        self.lr = arg.lr\n",
    "        self.dropout_rate = arg.dropout_rate\n",
    "        \n",
    "        # input placeholder\n",
    "        self.x = tf.placeholder(tf.int32, shape=(None, None))\n",
    "        self.y = tf.placeholder(tf.int32, shape=(None, None))\n",
    "        self.de_inp = tf.placeholder(tf.int32, shape=(None, None))\n",
    "        \n",
    "        # Encoder\n",
    "        with tf.variable_scope(\"encoder\"):\n",
    "            # embedding\n",
    "            self.en_emb = embedding(self.x, vocab_size=self.input_vocab_size, num_units=self.hidden_units, scale=True, scope=\"enc_embed\")\n",
    "            self.enc = self.en_emb + embedding(tf.tile(tf.expand_dims(tf.range(tf.shape(self.x)[1]), 0), [tf.shape(self.x)[0], 1]),\n",
    "                                          vocab_size=self.max_length,num_units=self.hidden_units, zero_pad=False, scale=False,scope=\"enc_pe\")\n",
    "            ## Dropout\n",
    "            self.enc = tf.layers.dropout(self.enc, \n",
    "                                        rate=self.dropout_rate, \n",
    "                                        training=tf.convert_to_tensor(self.is_training))\n",
    "\n",
    "            ## Blocks\n",
    "            for i in range(self.num_blocks):\n",
    "                with tf.variable_scope(\"num_blocks_{}\".format(i)):\n",
    "                    ### Multihead Attention\n",
    "                    self.enc = multihead_attention(key_emb = self.en_emb,\n",
    "                                                   que_emb = self.en_emb,\n",
    "                                                   queries=self.enc, \n",
    "                                                    keys=self.enc, \n",
    "                                                    num_units=self.hidden_units, \n",
    "                                                    num_heads=self.num_heads, \n",
    "                                                    dropout_rate=self.dropout_rate,\n",
    "                                                    is_training=self.is_training,\n",
    "                                                    causality=False)\n",
    "\n",
    "            ### Feed Forward\n",
    "            self.enc = feedforward(self.enc, num_units=[4*self.hidden_units, self.hidden_units])\n",
    "        \n",
    "        # Decoder\n",
    "        with tf.variable_scope(\"decoder\"):\n",
    "            # embedding\n",
    "            self.de_emb = embedding(self.de_inp, vocab_size=self.label_vocab_size, num_units=self.hidden_units, scale=True, scope=\"dec_embed\")\n",
    "            self.dec = self.de_emb + embedding(tf.tile(tf.expand_dims(tf.range(tf.shape(self.de_inp)[1]), 0), [tf.shape(self.de_inp)[0], 1]),\n",
    "                                          vocab_size=self.max_length,num_units=self.hidden_units, zero_pad=False, scale=False,scope=\"dec_pe\")\n",
    "            ## Dropout\n",
    "            self.dec = tf.layers.dropout(self.dec, \n",
    "                                        rate=self.dropout_rate, \n",
    "                                        training=tf.convert_to_tensor(self.is_training))        \n",
    "\n",
    "            ## Multihead Attention ( self-attention)\n",
    "            for i in range(self.num_blocks):\n",
    "                with tf.variable_scope(\"num_blocks_{}\".format(i)):\n",
    "                    ### Multihead Attention\n",
    "                    self.dec = multihead_attention(key_emb = self.de_emb,\n",
    "                                                   que_emb = self.de_emb,\n",
    "                                                   queries=self.dec, \n",
    "                                                    keys=self.dec, \n",
    "                                                    num_units=self.hidden_units, \n",
    "                                                    num_heads=self.num_heads, \n",
    "                                                    dropout_rate=self.dropout_rate,\n",
    "                                                    is_training=self.is_training,\n",
    "                                                    causality=True,\n",
    "                                                    scope='self_attention')\n",
    "\n",
    "            ## Multihead Attention ( vanilla attention)\n",
    "            for i in range(self.num_blocks):\n",
    "                with tf.variable_scope(\"num_blocks_{}\".format(i)):\n",
    "                    ### Multihead Attention\n",
    "                    self.dec = multihead_attention(key_emb = self.en_emb,\n",
    "                                                   que_emb = self.de_emb,\n",
    "                                                   queries=self.dec, \n",
    "                                                    keys=self.enc, \n",
    "                                                    num_units=self.hidden_units, \n",
    "                                                    num_heads=self.num_heads, \n",
    "                                                    dropout_rate=self.dropout_rate,\n",
    "                                                    is_training=self.is_training,\n",
    "                                                    causality=True,\n",
    "                                                    scope='vanilla_attention') \n",
    "\n",
    "            ### Feed Forward\n",
    "            self.outputs = feedforward(self.dec, num_units=[4*self.hidden_units, self.hidden_units])\n",
    "                \n",
    "        # Final linear projection\n",
    "        self.logits = tf.layers.dense(self.outputs, self.label_vocab_size)\n",
    "        self.preds = tf.to_int32(tf.argmax(self.logits, axis=-1))\n",
    "        self.istarget = tf.to_float(tf.not_equal(self.y, 0))\n",
    "        self.acc = tf.reduce_sum(tf.to_float(tf.equal(self.preds, self.y))*self.istarget)/ (tf.reduce_sum(self.istarget))\n",
    "        tf.summary.scalar('acc', self.acc)\n",
    "                \n",
    "        if is_training:  \n",
    "            # Loss\n",
    "            self.y_smoothed = label_smoothing(tf.one_hot(self.y, depth=self.label_vocab_size))\n",
    "            self.loss = tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.logits, labels=self.y_smoothed)\n",
    "            self.mean_loss = tf.reduce_sum(self.loss*self.istarget) / (tf.reduce_sum(self.istarget))\n",
    "               \n",
    "            # Training Scheme\n",
    "            self.global_step = tf.Variable(0, name='global_step', trainable=False)\n",
    "            self.optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=0.9, beta2=0.98, epsilon=1e-8)\n",
    "            self.train_op = self.optimizer.minimize(self.mean_loss, global_step=self.global_step)\n",
    "                   \n",
    "            # Summary \n",
    "            tf.summary.scalar('mean_loss', self.mean_loss)\n",
    "            self.merged = tf.summary.merge_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gO79HjmoGBXM"
   },
   "source": [
    "## 3. 训练模型\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "upKSv_ylJ_tx"
   },
   "source": [
    "### 3.1 参数设定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "twQC0Km6GBXN"
   },
   "outputs": [],
   "source": [
    "def create_hparams():\n",
    "    params = tf.contrib.training.HParams(\n",
    "        num_heads = 8,\n",
    "        num_blocks = 6,\n",
    "        # vocab\n",
    "        input_vocab_size = 50,\n",
    "        label_vocab_size = 50,\n",
    "        # embedding size\n",
    "        max_length = 100,\n",
    "        hidden_units = 512,\n",
    "        dropout_rate = 0.2,\n",
    "        lr = 0.0003,\n",
    "        is_training = True)\n",
    "    return params\n",
    "\n",
    "        \n",
    "arg = create_hparams()\n",
    "arg.input_vocab_size = len(encoder_vocab)\n",
    "arg.label_vocab_size = len(decoder_vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OO0OJOYwGBXP"
   },
   "source": [
    "### 3.2 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 557
    },
    "colab_type": "code",
    "id": "QHmUtPZ4GBXR",
    "outputId": "9ca91084-e3c8-4908-a395-d34e5b995e69"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 156/156 [00:31<00:00,  6.19it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  5.83it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.23it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.11it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.14it/s]\n",
      "  1%|          | 1/156 [00:00<00:27,  5.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epochs 5 : average loss =  3.3463863134384155\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 156/156 [00:23<00:00,  6.27it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  5.86it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  6.33it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.08it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  6.29it/s]\n",
      "  1%|          | 1/156 [00:00<00:26,  5.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epochs 10 : average loss =  2.0142565186207113\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 156/156 [00:24<00:00,  6.18it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  5.84it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.10it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.10it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  6.38it/s]\n",
      "  1%|          | 1/156 [00:00<00:27,  5.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epochs 15 : average loss =  1.5278632457439716\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 156/156 [00:24<00:00,  6.15it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  5.86it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.23it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  6.13it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  6.32it/s]\n",
      "  1%|          | 1/156 [00:00<00:27,  5.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epochs 20 : average loss =  1.4216684783116365\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 156/156 [00:23<00:00,  6.26it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  5.89it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.26it/s]\n",
      "100%|██████████| 156/156 [00:24<00:00,  6.10it/s]\n",
      "100%|██████████| 156/156 [00:23<00:00,  6.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epochs 25 : average loss =  1.3833287457625072\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "epochs = 25\n",
    "batch_size = 64\n",
    "\n",
    "g = Graph(arg)\n",
    "\n",
    "saver =tf.train.Saver()\n",
    "with tf.Session() as sess:\n",
    "    merged = tf.summary.merge_all()\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    if os.path.exists('logs/model.meta'):\n",
    "        saver.restore(sess, 'logs/model')\n",
    "    writer = tf.summary.FileWriter('tensorboard/lm', tf.get_default_graph())\n",
    "    for k in range(epochs):\n",
    "        total_loss = 0\n",
    "        batch_num = len(encoder_inputs) // batch_size\n",
    "        batch = get_batch(encoder_inputs, decoder_inputs, decoder_targets, batch_size)\n",
    "        for i in tqdm(range(batch_num)):\n",
    "            encoder_input, decoder_input, decoder_target = next(batch)\n",
    "            feed = {g.x: encoder_input, g.y: decoder_target, g.de_inp:decoder_input}\n",
    "            cost,_ = sess.run([g.mean_loss,g.train_op], feed_dict=feed)\n",
    "            total_loss += cost\n",
    "            if (k * batch_num + i) % 10 == 0:\n",
    "                rs=sess.run(merged, feed_dict=feed)\n",
    "                writer.add_summary(rs, k * batch_num + i)\n",
    "        if (k+1) % 5 == 0:\n",
    "            print('epochs', k+1, ': average loss = ', total_loss/batch_num)\n",
    "    saver.save(sess, 'logs/model')\n",
    "    writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "C4JtOyPoGBXW"
   },
   "source": [
    "### 3.3 模型推断"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 485
    },
    "colab_type": "code",
    "id": "6yHefDhGGBXY",
    "outputId": "c75eb8a1-60c4-48f6-f7a7-f057cb081709"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from logs/model\n",
      "输入测试拼音: You could be right, I suppose\n",
      "我猜想你可能是对的\n",
      "输入测试拼音: You don't believe Tom, do you\n",
      "你不信任汤姆，对吗\n",
      "输入测试拼音: Tom has lived here since 2003\n",
      "汤姆自从2003年就住在这里\n",
      "输入测试拼音: Tom asked if I'd found my key\n",
      "湯姆問我找到我的鑰匙了吗\n",
      "输入测试拼音: They have a very nice veranda\n",
      "他们有一个非常漂亮的暖房\n",
      "输入测试拼音: She was married to a rich man\n",
      "她嫁給了一個有錢的男人\n",
      "输入测试拼音: My parents sent me a postcard\n",
      "我父母給我寄了一張明信片\n",
      "输入测试拼音: Just put yourself in my shoes\n",
      "你站在我的立場上考慮看看\n",
      "输入测试拼音: It was a very stupid decision\n",
      "这是一个十分愚蠢的决定\n",
      "输入测试拼音: I'm really sorry to hear that\n",
      "听到这样的消息我真的很难过\n",
      "输入测试拼音: His wife is one of my friends\n",
      "他的妻子是我的一個朋友\n",
      "输入测试拼音: He thought of a good solution\n",
      "他想到了一個解決的好辦法\n",
      "输入测试拼音: exit\n"
     ]
    }
   ],
   "source": [
    "arg.is_training = False\n",
    "\n",
    "g = Graph(arg)\n",
    "\n",
    "saver =tf.train.Saver()\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    saver.restore(sess, 'logs/model')\n",
    "    while True:\n",
    "        line = input('输入测试拼音: ')\n",
    "        if line == 'exit': break\n",
    "        line = line.lower().replace(',', ' ,').strip('\\n').split(' ')\n",
    "        x = np.array([encoder_vocab.index(pny) for pny in line])\n",
    "        x = x.reshape(1, -1)\n",
    "        de_inp = [[decoder_vocab.index('<GO>')]]\n",
    "        while True:\n",
    "            y = np.array(de_inp)\n",
    "            preds = sess.run(g.preds, {g.x: x, g.de_inp: y})\n",
    "            if preds[0][-1] == decoder_vocab.index('<EOS>'):\n",
    "                break\n",
    "            de_inp[0].append(preds[0][-1])\n",
    "        got = ''.join(decoder_vocab[idx] for idx in de_inp[0][1:])\n",
    "        print(got)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qql1_AOlQnog"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "K0REVxb6GBV0",
    "_FDZYnOuGBW3",
    "IlpHtw43GBW7",
    "iwSIPvYtGBXK",
    "upKSv_ylJ_tx"
   ],
   "name": "AttentionIsAllYouNeed.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
